package org.github.manx98.jlibhac.tools.fssystem;

import org.github.manx98.jlibhac.common.Utilities;
import org.github.manx98.jlibhac.common.math.UnsignedNumber;
import org.github.manx98.jlibhac.common.math.Uint16;
import org.github.manx98.jlibhac.common.math.Uint64;
import org.github.manx98.jlibhac.util.BitUtil;

import javax.crypto.Cipher;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.ShortBufferException;
import javax.crypto.spec.SecretKeySpec;
import java.nio.ByteOrder;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;

/**
 * @author Manx98
 * @version 1.0
 * @since 2022/10/3 22:18
 */
public class Aes128CtrTransform {
    private static final int BlockSize = 128;
    private static final int BlockSizeBytes = BlockSize / 8;
    public final byte[] Counter = new byte[BlockSizeBytes];
    private final Cipher encryptor;

    public Aes128CtrTransform(byte[] key, byte[] counter) {
        if (key == null) throw new IllegalArgumentException("key cannot be null");
        if (counter == null) throw new IllegalArgumentException("counter cannot be null");
        if (key.length != BlockSizeBytes)
            throw new IllegalArgumentException("key must be " + BlockSizeBytes + " bytes long");
        if (counter.length != BlockSizeBytes)
            throw new IllegalArgumentException("counter must be " + BlockSizeBytes + " bytes long");
        try {
            this.encryptor = Cipher.getInstance("AES/ECB/NoPadding");
            encryptor.init(Cipher.ENCRYPT_MODE, new SecretKeySpec(key, "AES"));
        } catch (NoSuchAlgorithmException | NoSuchPaddingException | InvalidKeyException e) {
            throw new RuntimeException(e);
        }
        System.arraycopy(counter, 0, Counter, 0, BlockSizeBytes);
    }

    public void transformBlock(byte[] data, int offset, int size) {
        int blockCount = (int) BitUtil.divideUp(data.length, BlockSizeBytes);
        int length = blockCount * BlockSizeBytes;

        byte[] counterXor = new byte[length];
        System.arraycopy(Counter, 0, counterXor, 0, Counter.length);
        fillDecryptedCounter(counterXor);
        try {
            encryptor.update(counterXor, 0, counterXor.length, counterXor, 0);
        } catch (ShortBufferException e) {
            throw new RuntimeException(e);
        }
        Utilities.xorArrays(data, offset, size, counterXor);
    }

    public static void fillDecryptedCounter(byte[] buffer) {
        UnsignedNumber hi = new Uint64(buffer);
        UnsignedNumber lo = new Uint64(buffer, 8, ByteOrder.BIG_ENDIAN);
        UnsignedNumber one = new Uint16(1);
        for (int i = 2; i < buffer.length / 8; i += 2) {
            lo = lo.add(one);
            hi.toByteArray(buffer, i * 8, ByteOrder.LITTLE_ENDIAN);
            lo.toByteArray(buffer, (i + 1) * 8, ByteOrder.BIG_ENDIAN);
        }
    }
}
